'''
Created on Jul 16, 2009

@author: mkiyer
'''
import re
import sys

hs_ref_re = re.compile(r'hs_ref_(\w+).fa')

def __get_num_mismatches(match_descriptor):
    '''
    Given a match descriptor string, return the number of mismatching bases
    
    >>> __get_num_mismatches('11A')
    1
    >>> __get_num_mismatches('12T10G8TTT')
    5    
    '''
    # TODO: num_mismatches is time consuming -- optimize this       
    num_mismatches = 0
    for c in match_descriptor:
        if c in frozenset(['A', 'T', 'G', 'C']):
            num_mismatches += 1
    return num_mismatches

class ParseError(Exception):
    """Exception raised for errors parsing input
    
    Attributes: message -- explanation of the error
    """
    def __init__(self, msg):
        self.msg = msg
    def __str__(self):
        return self.msg

class ElandMapping(object):
    '''Describes the location of a sequence alignment
    
    Attributes:
    chrom - chromosome name of the alignment in string format 'chrN'
    chrom_start - integer location of start of alignment based on (+) strand
    chrom_end - integer location of end of alignment
    num_mismatches - number of alignment mismatches
    '''
    def __init__(self, chrom, chrom_start, chrom_end, strand, num_mismatches=0):
        self.chrom = chrom
        self.chrom_start= chrom_start
        self.chrom_end = chrom_end
        self.strand = strand
        self.num_mismatches = num_mismatches
    
    def __repr__(self):
        return self.__class__.__name__ + \
            "('%s', '%s', '%s', '%s', '%s')" % \
            (self.chrom, self.chrom_start, self.chrom_end, self.strand, 
             self.num_mismatches)
    
    def maps_to_genome(self):
        return self.chrom.startswith('chr')
    def maps_to_genes(self):
        return self.chrom.startswith('ucsc')
    
class ElandParser(object):
    def __init__(self):
        pass

    def parse_file(self, thisfhd):
        i = 0
        m = 0
        for thisline in thisfhd:
            # count lines for status updates
            i+=1
            if i == 100000:
                m += 1
                sys.stderr.write(" %d\n" % (m*100000))
                i=0
            # process line
            mappings = self.parse_line(thisline)
            if mappings:
                yield mappings

class ElandResultParser(ElandParser):
    '''
    File parser for Eland Result file
    
    >PATHBIO-SOLEXA_20E6WAAXX:8:1:914:878   \
        GGCTGATAGTGTCAGATTTCGTGTTGATTTTTTTTT    U2      0       0       1 \
        hs_ref_chr6.fa  44343560        R       ..      8C      30T

    Fields:
    0. Sequence name (derived from file name and line number if format is not fasta) 
    1. Sequence 
    2. Type of match codes:
        NM-No match found
        QC-No matching done: QC failure (too many Ns) 
        RM-No matching done: repeat masked (may be seen if repeatFile.txt was specified) 
        U0-Best match found was a unique exact match 
        U1-Best match found was a unique 1-error match 
        U2-Best match found was a unique 2-error match 
        R0-Multiple exact matches found 
        R1-Multiple 1-error matches found, no exact matches 
        R2-Multiple 2-error matches found, no exact or 1-error matches 
    3. Number of exact matches found 
    4. Number of 1-error matches found 
    5. Number of 2-error matches found 
    The following fields are only used if a unique best match was found: 
    6. Genome file in which match was found 
    7. Position of match (bases in file are numbered starting at 1) 
    8.Direction of match (F=forward strand, R=reverse) 
    9.How N characters in read were interpreted ("."=not applicable, "D"=Detection, "I"=Insertion) 
    The following field is only used in the case of a unique inexact match: 
    10.Position and type of first substitution error (A numeral refers to a run of matching 
    bases, an upper case base or N refers to a base in the reference that differs from the 
    read. For example, 11A: after 11 matching bases, base 12 is A in the reference but 
    not in the read)
    '''
    def __init__(self):
        super(ElandResultParser, self).__init__()
        
    def parse_line(self, thisline):
        # check to make sure line can be parsed
        if not thisline:
            raise ParseError("line is None")
        thisline = thisline.strip()
        if not thisline:
            raise ParseError("blank line")
        if thisline.startswith("#"):
            # comment line
            return None

        fields = thisline.split('\t')
        seqlen = len(fields[1])
        match_code = fields[2]
    
        if match_code == "U0" or match_code == "U1" or match_code == "U2":
            chrom = fields[6]
            if chrom.startswith('hs_ref'):               
                # convert chrom from eland file name to convention name
                chrom = hs_ref_re.match(chrom).group(1)              
            chrom_start = int(fields[7])-1
            strand = fields[8]
            #num_mismatches = __get_num_mismatches(match_descriptor):
            num_mismatches = None
            return [ElandMapping(chrom=chrom,
                                 chrom_start=chrom_start,
                                 chrom_end=chrom_start + seqlen,
                                 strand=strand,
                                 num_mismatches=num_mismatches)]
        else:
            # non-mapping
            return None

class ElandExportParser(ElandParser):
    '''
    parse lines of the eland export file and returns ElandMapping objects
    
    File format is tab-delimited fields as specified in the Solexa Pipeline documentation:
    0. machine
    1. run_number
    2. lane
    3. tile
    4. x coordinate of cluster
    5. y coordinate of cluster
    6. index string (bland for a non-indexed run)
    7. read number (1 or 2 for paired-read analysis, blank for a single-read analysis)
    8. read sequence
    9. quality string
    10. match chromosome (or code describing why no match occurred)
    11. match contig
    12. match position (always with respect to forward strand, eland numbering starts at one)
    13. match strand ('F' for forward, 'R' for reverse)
    14. match descriptor
    15. single-read alignment score
    16. paired-read alignment score
    17. partner chromosome
    18. partner contig
    19. partner offset
    20. partner strand
    22. filtering ('Y' for yes, 'N' for no)
    '''
    def __init__(self):
        super(ElandExportParser, self).__init__()

    def parse_line(self, thisline):
        # check to make sure line can be parsed
        if not thisline:
            raise ParseError("line is None")
        thisline = thisline.strip()
        if not thisline:
            raise ParseError("blank line")
        if thisline.startswith("#"):
            return None
        # split the line into fields and extract relevant information
        fields = thisline.split('\t')
        seqlen = len(fields[8])
        chrom = fields[10]
        
        if chrom.startswith('hs_ref') or chrom.startswith('ucsc'):
            # convert chrom from eland file name to convention name
            if chrom.startswith('hs_ref'):
                chrom = hs_ref_re.match(chrom).group(1)
            #chrom_start = (int(fields[12])-1) if fields[12] else 0
            chrom_start = int(fields[12]) - 1
            strand = fields[13]
            #num_mismatches = __get_num_mismatches(fields[14])
            num_mismatches = None
            return [ElandMapping(chrom=chrom,
                                 chrom_start=chrom_start,
                                 chrom_end=chrom_start + seqlen,
                                 strand=strand,
                                 num_mismatches=num_mismatches)]
        else:
            # non-mapping
            return None

class ElandMultiParser(ElandParser):
    """File Parser Class for ELAND multi File.

    Note this parser can only work for s_N_eland_multi.txt format.

    Each line of the output file contains the following fields: 
    1. Sequence name 
    2. Sequence 
    3. Either NM, QC, RM (as described above) or the following: 
    4. x:y:z where x, y, and z are the number of exact, single-error, and 2-error matches 
    found 
    5. Blank, if no matches found or if too many matches found, or the following: 
    BAC_plus_vector.fa:163022R1,170128F2,E_coli.fa:3909847R1 
    This says there are two matches to BAC_plus_vector.fa: one in the reverse direction 
    starting at position 160322 with one error, one in the forward direction starting at 
    position 170128 with two errors. There is also a single-error match to E_coli.fa.
    """
    def __init__(self):
        super(ElandMultiParser, self).__init__()

    def get_match_counts(self):
        self.match_code = map(int, self.match_code.split(':'))

    def parse_line(self, thisline):
        # check to make sure line can be parsed
        if not thisline:
            raise ParseError("line is None")
        thisline = thisline.strip()
        if not thisline:
            raise ParseError("blank line")
        if thisline.startswith("#"):
            # comment line
            return None
        
        fields = thisline.split('\t')
        if len(fields) < 4:
            # non-mapping
            return None
        
        name = fields[0]
        seqlen = len(fields[1])
        match_code = fields[2]
        matches = []
        chrom = None
        # extract all the matches
        for matchstring in fields[3].split(','):
            # check if contig (chromosome) is present in this string,
            # and if so suck it out of the matchstring
            matchstring = matchstring.split(':')
            if len(matchstring) == 2:
                chrom, matchstring = matchstring
                if chrom.startswith('hs_ref'):               
                    # convert chrom from eland file name to convention name
                    chrom = hs_ref_re.match(chrom).group(1)                
            else:
                matchstring = matchstring[0]
            # separate strand and number of mismatches from position
            # fortunately the number of mismatches is guaranteed to be
            # 0, 1, or 2, so it will always consume just one character
            strand = matchstring[-2]
            chrom_start = int(matchstring[:-2]) - 1
            chrom_end = chrom_start + seqlen
            num_mismatches = matchstring[-1]
            # append a new alignment to the list of matches
            matches.append(ElandMapping(chrom, chrom_start, 
                                        chrom_start + seqlen,
                                        strand, num_mismatches))
        return matches